package com.koro.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

@Tag(name = "Ollama AI")
@RestController
@RequestMapping("/ai")
public class OllamaController {

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private OllamaChatModel chatModel;

    @GetMapping("/chatModel")
    @Operation(summary = "智能问答")
    public String chatModel(@RequestParam(value = "message",defaultValue = "给我讲个笑话")String message) {
        ChatResponse response = chatModel.call(
                new Prompt(
                        message,
                        OllamaOptions.create()
                                .withModel(String.valueOf(OllamaModel.LLAMA3))
                                .withTemperature(0.4F)
        ));
        return response.getResult().getOutput().getContent();
    }

    @GetMapping(value = "/stream")
    @Operation(summary = "智能问答（流式）")
    public Flux<ChatResponse> chatStream(@RequestParam(value = "message",defaultValue = "给我讲个笑话")String message) {
        Flux<ChatResponse> stream = chatModel.stream(
                new Prompt(
                        message,
                        OllamaOptions.create()
                                .withModel(String.valueOf(OllamaModel.LLAMA3))
                                .withTemperature(0.4F)
                ));
        return stream;

    }
}
